# -*-coding:utf-8 -*-

"""
# File      : pdf_picture_extract.py
# Time      : 2023/9/9:15:36
# Author    : Black Curry
# Email     : 1832647481@qq.com
# Description : 从pdf中提取图片
# 对外提供接口: pdf2pic_with_page(path, fig_path)
"""

import fitz
import os
import time
import re
from math import fabs
import logging
# from YOLO_picextract import pdf2bigimg

logger = logging.getLogger(__name__)

# 设置图像尺寸的最小值
min_width = 100
min_height = 100
EPS = 1e-5
ERR_THRSHOLD = 50000

# logo_list_path = file_path = os.path.join(os.path.dirname(__file__), 'logo_list.txt')
# print(logo_list_path)
# with open(logo_list_path, 'r') as f:
#     lines = f.readlines()
#     logo_list = [line.strip() for line in lines]

# def check_logo(name, logo_list):
#     for logo in logo_list:
#         logo_data = logo.split("^")
#         if fabs(name[1] - float(logo_data[1])) < 20 and fabs(name[2] - float(logo_data[2])) < 20 and \
#             fabs(name[3] - float(logo_data[3])) < 20 and  fabs(name[4] - float(logo_data[4])) < 20 and\
#             fabs(name[5] - float(logo_data[5])) < 20 and fabs(name[6] - float(logo_data[6])) < 20:
#             return False

#     return True
"""
函数说明: 获取正确的保存路径；检查如果保存路径存在，则删除该路径下的所有文件，如果不存在，则创建该路径
"""


def get_save_path(fig_path):
    save_path = fig_path
    if os.path.exists(save_path):
        fig_list = os.listdir(save_path)
        for fig in fig_list:
            # print(fig)
            tem = os.path.join(save_path, fig)
            os.remove(tem)
    else:
        os.makedirs(save_path, exist_ok = True)
    return save_path


"""
函数说明: pdf提取图片的核心代码
"""
# def extract_picture(path, save_path):
#     image_meta_info = pdf2bigimg(path, save_path)
#     return image_meta_info

def extract_picture(path, save_path, checkIM):
    # 从文件名中提取paper_id
    paper_id = os.path.basename(path).split('.')[0]

    try:
        doc = fitz.open(path)
        page_nums = len(doc)
    except Exception as e:
        logger.exception("Exception occurred")
        raise Exception(f'无法打开 {path}')
    
    print("文件名:{}, 页数: {}".format(path, page_nums))

    docimg_list = []
    name_list = []
    err_count = 0
    fig_index = 1  # 将fig_index移到这里定义
    
    for page_num in range(page_nums):
        pageimg_list = []
        page = doc[page_num]

        image_bbox = doc.get_page_images(page_num)
        print(f"第{page_num + 1}页找到{len(image_bbox)}个图像对象")  # 调试信息

        for img in image_bbox:
            print(f"图像对象信息: {img}")  # 调试信息
            xref = img[0]
            
            # 修改检测逻辑：检查是否有有效的宽度和高度，而不是依赖正则表达式
            if len(img) >= 5 and img[2] > 0 and img[3] > 0:  # img[2]是宽度，img[3]是高度
                isImage = True
                print("检测到有效图像对象")
            else:
                isImage = False
                print("不是有效的图像对象，跳过")
                continue

            try:
                pix = fitz.Pixmap(doc, xref)
                
                # 如果不是RGB格式，转换为RGB
                if pix.colorspace.name != fitz.csRGB.name:
                    pix = fitz.Pixmap(fitz.csRGB, pix)

                width, height = pix.width, pix.height
                w = abs(width)
                h = abs(height)
                x0 = float(img[1])
                y0 = float(img[2])

                print(f"图像尺寸: {w}x{h}, 最小要求: {min_width}x{min_height}")  # 调试信息

                # 跳过不符合条件的图像
                if w < min_width or h < min_height or width * height < min_width * min_height:
                    print(f"图像尺寸不符合要求，跳过")  # 调试信息
                    err_count += 1
                    continue

                pix0 = fitz.Pixmap(fitz.csRGB, pix)
                new_name = f"{paper_id}_{page_num + 1}_{fig_index}.png"
                if new_name in name_list:
                    print(f"文件名重复，跳过: {new_name}")  # 调试信息
                    continue
                else:
                    print(f"成功添加图像: {new_name}")  # 调试信息
                    name_list.append(new_name)
                    pageimg_list.append([page_num + 1, w, h, x0, y0, width, height, pix0])
                    fig_index += 1  # 在成功添加图片后递增fig_index

            except Exception as e:
                print(f"处理图像时出错: {e}")  # 调试信息
                logger.exception("Exception occurred")

        docimg_list.append(pageimg_list)

    print('err_count :', err_count)
    if err_count > ERR_THRSHOLD:
        raise Exception(f'错误共出现 {err_count} 次 , 请人工检查{path} 源文件')

    image_meta_info = []
    for index_page, pageimg in enumerate(docimg_list):
        for index_img, img in enumerate(pageimg):
            # 重新生成文件名和保存图片
            new_name = f"{paper_id}_{img[0]}_{index_img + 1}.png"
            fig_path = os.path.join(save_path, new_name)
            img[7].save(fig_path)  # 保存图片到磁盘
            
            res = {
                "page_no": img[0],
                "fig_path": fig_path,
                "fig_width": img[1],
                "fig_height": img[2],
                "fig_x0": img[3],
                "fig_y0": img[4],
                "fig_scale_x": img[5],
                "fig_scale_y": img[6],
            }
            image_meta_info.append(res)

    print(f"提取了{len(image_meta_info)}张图片")
    doc.close()
    return image_meta_info


"""
函数说明: 检查是否是跨页图片
"""


def check_overpage(index_page, index_img, docimg_list):
    if index_page == 0:
        return True
    if docimg_list[index_page - 1] == []:
        return True
    cur_img = docimg_list[index_page][index_img]
    for img in docimg_list[index_page-1]:
        if fabs(cur_img[1] - img[1]) < 1 and fabs(cur_img[2] - img[2]) < 1 and \
                fabs(cur_img[5] - img[5]) < 1 and fabs(cur_img[6] - img[6]) < 1:
                # fabs(cur_img[3] - img[3]) < 1 and \

            return False
    return True


"""
函数说明: 从pdf中提取图片，将提取到的图像保存到指定路径，图像文件名包含页码、坐标。长宽等信息
参数说明:
    path: pdf文件路径
    fig_path: 图像保存路径
"""


def pdf2pic_with_page(path, fig_path):
    save_path = get_save_path(fig_path)
    # 设置图像尺寸的最小值
    # 使用正则表达式来查找图片
    # checkXO = r"/Type(?= */XObject)"

    # pdf提取图片的核心代码
    checkIM = r"/Subtype(?= */Image)"
    image_meta_info = extract_picture(path, save_path, checkIM)

    # image_meta_info = extract_picture(path, save_path)
    return image_meta_info

if __name__ == "__main__":
    path = "/home/zhangbo/workspace/aigc/genome/base_dataset/scripts/output/pdf/2010072773/2010072773.pdf"
    fig_path = "/home/zhangbo/workspace/aigc/genome/base_dataset/scripts/output/pdf/2010072773/figure"
    # 删除目录
    import shutil
    if os.path.exists(fig_path):
        shutil.rmtree(fig_path)
    os.makedirs(fig_path, exist_ok=True)
    image_meta_info = pdf2pic_with_page(path, fig_path)
    print(image_meta_info)
